
import os
import torch
import copy
from nets.fused_mobilenet_v3 import mobilenet_v3

siamese_model_path = 'model_data/best_epoch_weights.pth'

save_path = 'siamese'
save_model_path = os.path.join(save_path, 'best_epoch_weights.pth')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

siamese_net_dict = torch.load(siamese_model_path, map_location=device)

net = mobilenet_v3(mode='small', num_classes=4)
net_dict = net.state_dict()
local_parameter = copy.deepcopy(net_dict)

for key, value in siamese_net_dict.items():
    print(key, ' ', value.shape)

# 遍历模型的参数
for model_name in net_dict:
    # 从siamese模型中寻找相关名字的参数
    siamese_para_key = 'mobilenet_v3.' + model_name
    if siamese_para_key in siamese_net_dict.keys():
        local_parameter[model_name] = siamese_net_dict[siamese_para_key]

if not os.path.exists(save_path):
    os.makedirs(save_path)
torch.save(local_parameter, save_model_path)


